import os
import numpy as np
import librosa
from sklearn.svm import SVC
from sklearn.ensemble import RandomForestClassifier, VotingClassifier
from sklearn.model_selection import train_test_split, GridSearchCV
from sklearn.metrics import accuracy_score
from sklearn.preprocessing import StandardScaler

def load_audio_librosa(file_path):
    try:
        audio, sample_rate = librosa.load(file_path, sr=None)
        return audio, sample_rate
    except Exception as e:
        print(f"Error loading {file_path}: {e}")
        return None, None

def extract_features(audio, sample_rate):
    mfcc = librosa.feature.mfcc(y=audio, sr=sample_rate, n_mfcc=40)
    mfcc_mean = np.mean(mfcc.T, axis=0)
    spectral_centroid = np.mean(librosa.feature.spectral_centroid(y=audio, sr=sample_rate))
    zero_crossing_rate = np.mean(librosa.feature.zero_crossing_rate(y=audio))
    chroma_stft = np.mean(librosa.feature.chroma_stft(y=audio, sr=sample_rate), axis=1)
    spectral_rolloff = np.mean(librosa.feature.spectral_rolloff(y=audio, sr=sample_rate))
    tonnetz = np.mean(librosa.feature.tonnetz(y=audio, sr=sample_rate), axis=1)
    chroma_cqt = np.mean(librosa.feature.chroma_cqt(y=audio, sr=sample_rate), axis=1)
    return np.hstack([mfcc_mean, spectral_centroid, zero_crossing_rate, chroma_stft, spectral_rolloff, tonnetz, chroma_cqt])

def prepare_features(data_folder):
    features, labels = [], []
    print(f"Checking directory: {data_folder}")
    if not os.path.exists(data_folder):
        print("Directory does not exist.")
        return np.array(features), np.array(labels)

    for subdir, dirs, files in os.walk(data_folder):
        category = os.path.basename(subdir)
        for file in files:
            if file.endswith(".ogg"):
                file_path = os.path.join(subdir, file)
                audio, sample_rate = load_audio_librosa(file_path)
                if audio is not None:
                    feature = extract_features(audio, sample_rate)
                    features.append(feature)
                    labels.append(category)
    return np.array(features), np.array(labels)

def train_svm_model(X_train, y_train, X_test, y_test):
    parameters = {'kernel':('linear', 'rbf'), 'C':[1, 10, 100, 1000], 'gamma':('scale', 'auto')}
    svc = SVC()
    clf = GridSearchCV(svc, parameters, cv=10)
    clf.fit(X_train, y_train)
    best_params = clf.best_params_
    accuracy = accuracy_score(y_test, clf.predict(X_test))
    print("Best parameters:", best_params)
    return clf, accuracy

def train_models(X, y):
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
    scaler = StandardScaler()
    X_train = scaler.fit_transform(X_train)
    X_test = scaler.transform(X_test)

    svm_model, svm_accuracy = train_svm_model(X_train, y_train, X_test, y_test)
    rf_model = RandomForestClassifier(n_estimators=200)
    rf_model.fit(X_train, y_train)
    rf_accuracy = accuracy_score(y_test, rf_model.predict(X_test))

    # Voting Classifier
    voting_clf = VotingClassifier(estimators=[
        ('svm', svm_model), ('rf', rf_model)], voting='hard')
    voting_clf.fit(X_train, y_train)
    voting_accuracy = accuracy_score(y_test, voting_clf.predict(X_test))

    return {
        'SVM Accuracy': svm_accuracy,
        'Random Forest Accuracy': rf_accuracy,
        'Voting Classifier Accuracy': voting_accuracy
    }

if __name__ == "__main__":
    data_folder = "data_folder"
    X, y = prepare_features(data_folder)
    print("Loaded features:", X.shape)
    print("Loaded labels:", y.shape)

    if X.shape[0] > 0:
        accuracies = train_models(X, y)
        print("Training results:", accuracies)
    else:
        print("No data loaded. Please check your data folder path and contents.")